import numpy as np
from matplotlib import pyplot as plt

from mnist.multi_layer_net_extend import MultiLayerNetExtend
from mnist_data import load_mnist
from optimizer import SGD,Adam
from multi_layer_net_extend import MultiLayerNetExtend

(x_train,t_train),(x_test,t_test)=load_mnist(normalize=True,one_hot_label=True)
#减少学习数据
x_train=x_train[:1000]
t_train=t_train[:1000]

max_epochs=20
train_size=x_train.shape[0]
batch_size=100
learning_rate=0.001

def __train(weight_init_std):
    #（1）初始化网络和优化器
    #使用了 batch_norm的
    bn_network=MultiLayerNetExtend(input_size=784,hidden_size_list=[100,100,100,100,100],
                             output_size=10,weight_init_std=weight_init_std,
                             use_batch_norm=True)
    #未使用batch_norm的
    network=MultiLayerNetExtend(input_size=784,hidden_size_list=[100,100,100,100,100],output_size=10,weight_init_std=weight_init_std)
    optimizer=SGD(lr=learning_rate)
    #（2）记录准确率
    train_acc_list=[]
    bn_train_acc_list=[]
    iter_per_epoch=max(train_size/batch_size,1)#每一轮迭代的次数
    epoch_cnt=0#训练轮次计数器

    for i in range(1000000000):
        batch_mask=np.random.choice(train_size,batch_size)
        x_batch=x_train[batch_mask]
        t_batch=t_train[batch_mask]
        for _network in (bn_network,network):
            grads=_network.gradient(x_batch,t_batch)
            optimizer.update(_network.params,grads)
        if i%iter_per_epoch==0:
            train_acc=network.accuracy(x_train,t_train)
            bn_test_acc=bn_network.accuracy(x_train,t_train)
            train_acc_list.append(train_acc)
            bn_train_acc_list.append(bn_test_acc)
            print("epoch:"+str(epoch_cnt)+"|"+str(train_acc)+"|"+str(bn_test_acc))

            epoch_cnt+=1
            if epoch_cnt>=max_epochs:
                break
    return train_acc_list,bn_train_acc_list

#开始绘制图形
#对比了不同权重初始化缩放系数（weight_scale)下，普通神经网络和带 batch_normalization的神经网络的训练准确率差异
#也就是说这个w是我们现在这里生成的这些随机的数字，作为参数的标准差（weight_init_std)
plt.figure(figsize=(12,12))
weight_scale_list=np.logspace(0,-4,num=16)#生成（1到 0.0001分布的共 16个值）
x=np.arange(max_epochs)
for i,w in enumerate(weight_scale_list):
    print("============"+str(i+1)+"/16"+"=============")#打印进度：1/16到 16/16
    train_acc_list,bn_train_acc_list=__train(w)
    #创建 4 行 4 列的子图，当前绘制到第 i+1个子图
    plt.subplot(4,4,i+1)
    plt.title("W:"+str(w))#子图标题：当前权重缩放系数（如 W:1.0;W:0.0001）
    #只有最后一个子图（i=15）显示图例，其他的不显示，避免混乱
    if i==15:
        #绘制batch_norm的准确率曲线：实线，每两个点标记一次
        plt.plot(x,bn_train_acc_list,label="Batch Normalization",markevery=2)
        plt.plot(x,train_acc_list,linestyle="--",label='Normal(without BatchNorm)',markevery=2)
    else:
        plt.plot(x,bn_train_acc_list,markevery=2)
        plt.plot(x,train_acc_list,linestyle="--",markevery=2)
    #y轴范围指定为 0～1.0
    plt.ylim(0,1.0)
    #控制y轴刻度，只有第1、5、9、13个子图（左列）显示 y 轴标签和刻度，其他子图隐藏 y 轴刻度（避免拥挤）
    if i%4:
        plt.yticks([])
    else:
        plt.ylabel("Accuracy")
    if i<12:
        plt.xticks([])
    else:
        plt.xlabel("Epoch")#只有最后一行子图显示x标签（训练轮次）
    plt.legend(loc="lower right")#显示图例，只有最后一个子图有效

plt.show()
